from locust import HttpUser, task, between
from utils.yaml_utils import read_api_yamls, read_config
from utils.token_utils import get_login_token
from utils.log_utils import logger
from utils.http_utils import HttpClient
from typing import Dict, Any

config = read_config()

# 获取当前激活的环境配置
active_env = config.get("active_env")
env_config = config.get("environments", {}).get(active_env, {})
host = env_config.get("host")

if not host:
    raise ValueError(f"未找到环境 {active_env} 的host配置")

try:
    token = get_login_token(config)
    if not token:
        raise ValueError("获取token失败")

    api_configs = read_api_yamls(token)
    all_api_configs: Dict[str, Any] = {
        config.get("api", {}).get("name", ""): config for config in api_configs
    }

except Exception as e:
    logger.error(f"初始化失败: {str(e)}")
    raise


class APITestUser(HttpUser):
    host = host
    wait_time = between(1, 3)  # 设置默认的等待时间范围

    @task
    def run_api_tests(self):
        # 每次任务执行时都重新读取 API 配置，确保数据的独立性
        api_configs = read_api_yamls(token)
        for api_config in api_configs:
            api = api_config.get("api")
            detail = api.get("detail")
            method = api.get("method").lower()
            path = api.get("path")
            headers = api.get("headers", {})
            data = api.get("data", {})
            assert_config = api.get("assert", {})

            url = f"{self.host}{path}"

            with self.client.request(
                method, url, headers=headers, json=data, catch_response=True
            ) as response:
                if assert_config:
                    HttpClient.assert_response(response, assert_config)
                else:
                    logger.warning(
                        "接口配置文件无 assert 字段，故无需进行断言，若需进行断言，请添加 assert 字段"
                    )

                logger.info("-" * 100)
                logger.info(f"接口名称：{detail}     API:{path}")
                logger.info(f"请求 URL: {url}")
                logger.info(f"请求头: {headers}")
                logger.info(f"请求体: {data}")
                logger.info(f"响应内容：{response.text}")
